import tushare as ts
from pymongo import MongoClient, DESCENDING, ASCENDING
import json
import datetime
import platform
print('python v', platform.python_version())

# 连接数据库
client = MongoClient('mongodb://localhost:27017')
database = client['stocks']

# 1. 获取所有股票列表
# stock_list_df = ts.get_stock_basics()


# 获取集合最近一条记录的时间
# 如果集合的date字段不是datetime类型则报错
def find_latest_date(stock):
    result = []
    if stock.count():
        latest = stock.find({}).sort('date', -1).limit(1)
        for x in latest:
            result.append(x)
    if len(result):
        return datetime.datetime.strptime(result[0].get('date').strftime('%Y%m%d') + ' 0:0:0', '%Y%m%d %H:%M:%S') + datetime.timedelta(1)
    else:
        return None

# 如果集合的最新日期不等于最近的一个交易日，就以最新日期为start，以当前日期为end获取当前股票历史数据，并且更新数据库，插入新记录

# dd/mm/yyyy格式
# print (time.strftime("%Y-%m-%d"))
# datetime.datetime.now()
count = 0
# 遍历股票列表并且存入数据库
for stock_info in database['stock_list'].find({}).__iter__():   # 获取每行的index、row
    item = stock_info
    if count < 2:
        count += 1
        # 2. 更新股票列表，stock_list集合使用股票代码作为索引
        stock_collection = database['stock_' + item.get('code')]
        stock_collection.create_index('date', unique=True)

        # 当前时间
        now = datetime.datetime.now()
        # 今天日期
        today = datetime.datetime.strptime(now.strftime('%Y%m%d') + ' 0:0:0', '%Y%m%d %H:%M:%S')
        # 默认日期
        default_date = today
        if item['exchange'] == 'SH':
            default_date = datetime.datetime.strptime('1990-11-26 0:0:0', '%Y-%m-%d %H:%M:%S')
        if item['exchange'] == 'SZ':
            default_date = datetime.datetime.strptime('1990-12-01 0:0:0', '%Y-%m-%d %H:%M:%S')
        # 更新数据的起始日期
        start_date = find_latest_date(stock_collection) or item['ipo_date'] or default_date
        # 今天收盘的时间
        today_close_date = datetime.datetime.strptime(now.strftime('%Y%m%d') + ' 15:0:0', '%Y%m%d %H:%M:%S')

        # 判断数据是否更新到最近日期，如果是则不再抓取，如果不是则更新
        if today == start_date or ((today - start_date).days == 1 and now < today_close_date):
            print(item["code"], '数据已是最新')
        else:
            start_date_str = start_date.strftime('%Y-%m-%d')
            end_date_str = today.strftime('%Y-%m-%d')
            print(item["code"], '获取日线数据', start_date_str, end_date_str)
            hist_data_df = ts.get_k_data(code=str(item.get('code')), start=start_date_str, end=end_date_str)
            if len(hist_data_df) > 0:
                daily_list = []
                for key, day_info in hist_data_df.iterrows():
                    day_info_dict = json.loads(day_info.to_json())
                    daily_list.append({
                        'date': datetime.datetime.strptime(str(day_info_dict['date']), '%Y-%m-%d'),
                        'open': day_info_dict['open'],
                        'close': day_info_dict['close'],
                        'high': day_info_dict['high'],
                        'low': day_info_dict['low'],
                        'volume': day_info_dict['volume'],
                        'amount': day_info_dict['amount']
                    })
                    print(day_info_dict)
                stock_collection.insert_many(daily_list)
                print(item["code"], '导入 ', start_date_str, '-', end_date_str, '共', len(hist_data_df), '条数据')
            else:
                print('可能由于停牌等原因，暂无最新数据')
